# 导入必要的库和模块
import gradio as gr
import pickle
import numpy as np
from PIL import Image

# 打印Gradio版本
print(f"Gradio version: {gr.__version__}")

# 加载保存的KNN模型
with open('best_knn_model.pkl', 'rb') as f:
    knn_model = pickle.load(f)

# 定义预处理函数
def preprocess_image(image):
    # 确保图像是PIL Image对象
    if not isinstance(image, Image.Image):
        image = Image.fromarray(image)
    
    # 将图像转换为灰度图
    gray_image = image.convert('L')
    # 调整图像大小为8x8像素
    resized_image = gray_image.resize((8, 8), Image.LANCZOS)
    # 将图像转换为numpy数组
    img_array = np.array(resized_image)
    # 将像素值归一化到0-16范围内
    img_array = img_array / 16.0
    # 将图像展平为一维数组
    return img_array.flatten()

# 定义预测函数
def predict_digit(image):
    preprocessed_image = preprocess_image(image)
    prediction = knn_model.predict([preprocessed_image])[0]
    return f"预测的数字是: {prediction}"

# 创建Gradio接口
iface = gr.Interface(
    fn=predict_digit,
    inputs=gr.Image(),
    outputs="text",
    title="手写数字识别",
    description="上传或绘制一个数字图像，模型将尝试识别它。"
)

# 启动Gradio接口
if __name__ == "__main__":
    iface.launch()